{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Grok-powered SBOM Analysis with AG2\n",
    "\n",
    "This notebook demonstrates how to use AG2 (AutoGen) with Grok LLM to perform Software Bill of Materials (SBOM) analysis and security vulnerability scanning on Git repositories.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "**Author:** Priyanshu Deshmukh"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "The notebook showcases:\n",
    "\n",
    "1. Setting up Grok LLM integration with AG2\n",
    "2. Creating specialized security agents for vulnerability scanning\n",
    "3. Using Trivy for automated SBOM analysis\n",
    "4. Implementing function calling for repository scanning\n",
    "5. Generating comprehensive security reports"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Requirements\n",
    "\n",
    "```bash\n",
    "pip install -q \"ag2[openai,rag]\"\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Ubuntu/Debian\n",
    "   ```bash\n",
    "   sudo apt-get install wget apt-transport-https gnupg lsb-release\n",
    "   wget -qO - https://aquasecurity.github.io/trivy-repo/deb/public.key | sudo apt-key add -\n",
    "   echo deb https://aquasecurity.github.io/trivy-repo/deb $(lsb_release -sc) main | sudo tee -a /etc/apt/sources.list.d/trivy.list\n",
    "   sudo apt-get update\n",
    "   sudo apt-get install trivy\n",
    "   ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "from autogen import AssistantAgent, ConversableAgent, LLMConfig\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "#### Configure Environment Variables\n",
    "   ```\n",
    "   XAI_API_KEY=\"your-grok-api-key\"\n",
    "   ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_config = LLMConfig(\n",
    "    {\n",
    "        \"model\": \"grok-4\",\n",
    "        \"api_type\": \"openai\",\n",
    "        \"base_url\": \"https://api.x.ai/v1\",\n",
    "        \"api_key\": os.getenv(\"XAI_API_KEY\"),\n",
    "    },\n",
    "    temperature=0.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "appsec_agent = ConversableAgent(\n",
    "    name=\"appsec_agent\",\n",
    "    system_message=\"You are a security expert who performs AppSec scans on git repositories.\",\n",
    "    human_input_mode=\"TERMINATE\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "executer = AssistantAgent(\n",
    "    name=\"executor_agent\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    llm_config=llm_config,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Function Definitions\n",
    "\n",
    "### `scan_git_repo_vulnerabilities`\n",
    "\n",
    "1. Performs automated vulnerability scanning on Git repositories using Trivy."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "#### Parameters:\n",
    "1. `git_repo_url` (str): The Git repository URL to scan  \n",
    "2. `output_json_path` (Optional[str]): Path to save JSON output"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "#### Returns:\n",
    "1. JSON string containing vulnerability scan results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "@appsec_agent.register_for_llm()\n",
    "@executer.register_for_execution()\n",
    "def scan_git_repo_vulnerabilities(\n",
    "    git_repo_url: str, output_json_path: str | None = \"../test/agentchat/contrib/graph_rag/scan.json\"\n",
    ") -> str:\n",
    "    \"\"\"\n",
    "    Execute Trivy repo scan on a git repository and save the JSON output to a file.\n",
    "\n",
    "    Args:\n",
    "        git_repo_url (str): The git repository URL to scan\n",
    "        output_json_path (Optional[str]): Path to save the JSON output. If None, does not save.\n",
    "\n",
    "    Returns:\n",
    "        str: The JSON output from Trivy scan as a string\n",
    "\n",
    "    Raises:\n",
    "        subprocess.CalledProcessError: If the Trivy command fails\n",
    "        FileNotFoundError: If Trivy is not installed or not in PATH\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # Execute the trivy command\n",
    "        result = subprocess.run(\n",
    "            [\"trivy\", \"repo\", \"--format=json\", git_repo_url], capture_output=True, text=True, check=True\n",
    "        )\n",
    "        json_output = result.stdout\n",
    "        if output_json_path:\n",
    "            with open(output_json_path, \"w\", encoding=\"utf-8\") as f:\n",
    "                f.write(json_output)\n",
    "        return f\"Scan result saved to {output_json_path} for : {git_repo_url} \\n {json_output}\"\n",
    "    except subprocess.CalledProcessError as e:\n",
    "        return f\"Error: {e.stderr}\"\n",
    "    except FileNotFoundError:\n",
    "        return \"Error: Trivy command not found. Please ensure Trivy is installed and in your PATH.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "appsec_agent.initiate_chat(\n",
    "    message=\"Scan the following git repository for vulnerabilities and give me an analysis report: https://github.com/ine-labs/AWSGoat.git\",\n",
    "    max_turns=5,\n",
    "    recipient=executer,\n",
    ").process()"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Using AG2 with Grok to perform SBOM analysis and security vulnerability scanning on Git repositories",
   "tags": [
    "grok",
    "openai"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
